use std::{
    collections::{btree_map::Range, BTreeMap},
    fs::{File, OpenOptions},
    io::{BufReader, BufWriter, Read, Result, Seek, Write},
    path::PathBuf,
};

type Db = BTreeMap<Vec<u8>, (i32, u32)>;

pub const LEN: u32 = 4;

pub trait DbOperate {
    fn put(&mut self, key: &Vec<u8>, value: Option<&Vec<u8>>) -> Result<()>;
    fn del(&mut self, key: &Vec<u8>) -> Result<()>;
    fn change(&mut self, key: &Vec<u8>, value: Option<&Vec<u8>>) -> Result<()>;
    fn get(&mut self, key: &Vec<u8>) -> Option<Vec<u8>>;
}

pub struct ScanIter<'a> {
    pub iter: Range<'a, Vec<u8>, (i32, u32)>,
    pub disk: &'a File,
}

impl<'a> Iterator for ScanIter<'a> {
    type Item = (Vec<u8>, Option<Vec<u8>>);

    fn next(&mut self) -> Option<<Self as Iterator>::Item> {
        self.iter.next().map(|(k, (val_len, pos))| {
            if *val_len == -1 {
                return ((*k).clone(), None);
            }

            let _ = self.disk.seek(std::io::SeekFrom::Start(*pos as u64));
            let mut value = vec![0u8; *val_len as usize];
            let _ = self.disk.read_exact(&mut value);

            ((*k).clone(), Some(value))
        })
    }
}

pub struct Storage {
    pub disk: File,
    pub cache: Db,
}

impl Storage {
    pub fn new(path_buf: PathBuf) -> Self {
        let file = match OpenOptions::new()
            .read(true)
            .write(true)
            .create(true) // 如果文件不存在则创建
            .open(&path_buf)
        {
            Ok(file) => file,
            Err(e) => {
                panic!("Error opening or creating file: {}", e);
            }
        };

        Storage {
            disk: file,
            cache: BTreeMap::new(),
        }
    }

    pub fn load(&mut self) -> Result<()> {
        let mut file_size = self.disk.metadata()?.len() as u32;
        let mut pos = 0;
        let mut read_buffer = BufReader::new(&mut self.disk);
        loop {
            if file_size <= 0 {
                break;
            }
            let mut key_len_buf: [u8; 4] = [0; 4];
            read_buffer.read_exact(&mut key_len_buf)?;
            let key_len = u32::from_be_bytes(key_len_buf);

            let mut value_len_buf: [u8; 4] = [0; 4];
            read_buffer.read_exact(&mut value_len_buf)?;
            let value_len = i32::from_be_bytes(value_len_buf);

            let mut key: Vec<u8> = vec![0; key_len as usize];
            read_buffer.read_exact(&mut key)?;

            let mut val_len: u32 = value_len as u32;
            if value_len == -1 {
                val_len = 0;
            }

            read_buffer.seek(std::io::SeekFrom::Current(val_len as i64))?;

            pos = pos + 4 * 2 + key_len;
            self.cache.insert(key, (value_len, pos));

            file_size = file_size - 4 - 4 - key_len - val_len;
            pos = pos + val_len;
        }
        Ok(())
    }

    pub fn scan(&self, range: impl std::ops::RangeBounds<Vec<u8>>) -> ScanIter<'_> {
        ScanIter {
            iter: self.cache.range(range),
            disk: &self.disk,
        }
    }
}

impl DbOperate for Storage {
    /**
     * | key len | value len | key | value |
     */
    fn put(&mut self, key: &Vec<u8>, value: Option<&Vec<u8>>) -> Result<()> {
        let key_len: u32 = key.len() as u32;
        let value_len = match value {
            None => -1,
            Some(value) => value.len() as i32,
        };

        let mut write_buffer = BufWriter::new(&mut self.disk);
        let result = write_buffer.seek(std::io::SeekFrom::End(0))?;

        let pos = LEN * 2 + key_len + result as u32;

        write_buffer.write_all(&key_len.to_be_bytes())?;
        write_buffer.write_all(&value_len.to_be_bytes())?;
        write_buffer.write_all(&key)?;
        if let Some(val) = value {
            write_buffer.write_all(&val)?;
        }
        write_buffer.flush()?;
        //println!("{:?}, {:?}", key.clone(), (value_len, pos));

        self.cache.insert(key.clone(), (value_len, pos));

        Ok(())
    }

    fn del(&mut self, key: &Vec<u8>) -> Result<()> {
        self.put(key, None)?;
        self.cache.remove(key);
        Ok(())
    }

    fn change(&mut self, key: &Vec<u8>, value: Option<&Vec<u8>>) -> Result<()> {
        self.put(key, value)
    }

    fn get(&mut self, key: &Vec<u8>) -> Option<Vec<u8>> {
        let pos = self.cache.get(key);
        match pos {
            None => None,
            Some((value_len, pos)) => {
                if *value_len == -1 {
                    return None;
                }
                let _ = self.disk.seek(std::io::SeekFrom::Start(*pos as u64));
                let mut value_buf: Vec<u8> = vec![0; *value_len as usize];
                let _ = self.disk.read_exact(&mut value_buf);
                Some(value_buf)
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use std::{io::Result, ops::Bound, path::PathBuf};

    use crate::{DbOperate, Storage};

    #[test]
    fn test_put() -> Result<()> {
        let path_buf = PathBuf::from("/tmp/kv_db/20240718");

        let mut storage = Storage::new(path_buf);
        storage.put(
            &"hello".as_bytes().to_vec(),
            Some(&"world".as_bytes().to_vec()),
        )?;
        storage.put(
            &"你好".as_bytes().to_vec(),
            Some(&"rust".as_bytes().to_vec()),
        )?;
        storage.put(&"none".as_bytes().to_vec(), None)?;
        storage.put(&"null".as_bytes().to_vec(), Some(&vec![]))?;
        storage.put(&"hi".as_bytes().to_vec(), Some(&"rust".as_bytes().to_vec()))?;

        let value = match storage.get(&"null".as_bytes().to_vec()) {
            Some(bytes) => String::from_utf8_lossy(&bytes).into_owned(),
            None => String::from("<not found>"), // 处理找不到的情况
        };
        println!("value = {}", value);

        Ok(())
    }

    #[test]
    fn test_load() -> Result<()> {
        let path_buf = PathBuf::from("/tmp/kv_db/20240718");

        let mut storage = Storage::new(path_buf);
        storage.load()?;

        let mut keys = Vec::<Vec<u8>>::new();
        for (k, _) in storage.cache.iter() {
            keys.push((*k).clone());
        }

        for key in keys.into_iter() {
            let value = match storage.get(&key) {
                Some(bytes) => String::from_utf8_lossy(&bytes).into_owned(),
                None => String::from("<not found>"), // 处理找不到的情况
            };
            println!("key = {}, value = {}", String::from_utf8_lossy(&key), value)
        }

        Ok(())
    }

    #[test]
    fn test_change() -> Result<()> {
        let path_buf = PathBuf::from("/tmp/kv_db/20240718");

        let mut storage = Storage::new(path_buf);
        storage.load()?;

        storage.change(
            &"null".as_bytes().to_vec(),
            Some(&"null".as_bytes().to_vec()),
        )?;

        let value = match storage.get(&"null".as_bytes().to_vec()) {
            Some(bytes) => String::from_utf8_lossy(&bytes).into_owned(),
            None => String::from("<not found>"), // 处理找不到的情况
        };
        println!("value = {}", value);

        Ok(())
    }

    #[test]
    fn test_del() -> Result<()> {
        let path_buf = PathBuf::from("/tmp/kv_db/20240718");

        let mut storage = Storage::new(path_buf);
        storage.load()?;

        storage.del(&"null".as_bytes().to_vec())?;

        let value = match storage.get(&"null".as_bytes().to_vec()) {
            Some(bytes) => String::from_utf8_lossy(&bytes).into_owned(),
            None => String::from("<not found>"), // 处理找不到的情况
        };
        println!("value = {}", value);

        Ok(())
    }

    #[test]
    fn test_scan() -> Result<()> {
        let path_buf = PathBuf::from("/tmp/kv_db/20240718");

        let mut storage = Storage::new(path_buf);
        storage.load()?;

        let start = Bound::Included("a".as_bytes().to_vec());
        let end = Bound::Excluded("z".as_bytes().to_vec());
        let scan_iter = storage.scan((start, end));

        for (k, v) in scan_iter {
            let value = match v {
                None => "not found".to_owned(),
                Some(value) => String::from_utf8_lossy(&value).into_owned(),
            };
            println!("key = {}, value = {}", String::from_utf8_lossy(&k), value);
        }
        Ok(())
    }
}
